import scipy.io as sio
import numpy as np
import os
from glob import glob
from matplotlib import pyplot as plt


if __name__ == "__main__":
    # 读取数据
    # mat_list = glob("SAFDNN/CRWU/12k Drive End Bearing Fault Data/**/*.mat", recursive=True)  # 训练数据
    mat_list = glob("SAFDNN/CRWU/12k Drive End Bearing Fault Data Test/**/*.mat", recursive=True)  # 测试数据
    stride = 1024
    repeat = 512
    arr = []
    for mat in mat_list:
        data = sio.loadmat(mat)
        if "Inner Race" in mat:
            label = 0
        elif "Outer Race" in mat:
            label = 1
        elif "Ball" in mat:
            label = 2
        else:
            label = 3
        for k in data.keys():
            if k.endswith("DE_time"):
                # ffted = abs(np.fft.fft(data[k].reshape(-1)))
                # x = np.fft.fftfreq(len(ffted), 1 / 12000)
                # plt.plot(x[:len(ffted) // 2], ffted[:len(ffted) // 2])
                # plt.vlines(300, 0, 1e3, colors="r", linestyles="dashed")
                # plt.show()
                child_data = data[k].reshape(-1)[10000:-10000]
                i = 0
                while stride+repeat*i <= len(child_data):
                    sample = child_data[repeat*i:stride+repeat*i]
                    sample = np.append(sample, label)
                    arr.append(sample)
                    i += 1
    arr = np.array(arr)
    np.save("SAFDNN/test_data_512.npy", arr)
